#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@author: liang kang
@contact: gangkanli1219@gmail.com
@time: 2018/2/28 17:17
@desc: 基本数据
"""
import os

from dltools.utils import log


class BaseFileReader(object):
    """
    基本原始图像数据读取
    """

    def __init__(self, root,
                 recurrence=True,
                 max_number=1e12,
                 display=100,
                 logger=None):
        """

        Parameters
        ----------
        root: 根目录
        recurrence: 是否递归搜索子目录
        max_number: 读入的最大样本数
        display: 显示间隔
        logger: 日志对象
        """
        self.root = root
        if not os.path.exists(self.root):
            raise NotADirectoryError

        # 当需要递归遍历目录时，该变量保存进入的子目录
        self.buf_root = root

        # 保存中间数据
        self.buf_data = None
        self.recurrence = recurrence

        self.max_number = max_number
        self.count = 0
        self.display = display

        if logger is not None:
            self.logger = logger
        else:
            self.logger = log.get_console_logger('FileReader')

    def __iter__(self):
        list_file = os.listdir(self.buf_root)
        for file in list_file:
            file_path = os.path.join(self.buf_root, file)
            if os.path.isfile(file_path):
                if self._filter(file_path):
                    if self.count % self.display == 0:
                        self.logger.info(
                            'Reading the number of {}.'.format(self.count))
                    yield self.buf_data
            elif self.recurrence:
                self.buf_root = file_path
                for result in self:
                    yield result
                self.buf_root = os.path.dirname(file_path)
            else:
                pass

    def has_next(self):
        """
        判断是否已经超过最大读取数量

        Returns
        -------

        """
        if self.count < self.max_number:
            self.count += 1
            return True
        else:
            raise StopIteration

    def _filter(self, file):
        """
        自定义用于筛选文件的方法, 可以重载

        Returns
        -------

        """
        if self.has_next():
            # 统一使用这样的方法表示数据
            self.buf_data = {'feature': file,  # feature 表示数据的特征
                             'label': None,  # label 表示数据的标记信息
                             'data': None,  # data 表示其他辅助信息
                             'name': file}  # name 表示数据的表示
            return True
        else:
            return False
